package org.hk.proxy.server.handlers;

import cn.hutool.core.util.ByteUtil;
import cn.hutool.core.util.StrUtil;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.handler.codec.socksx.v5.DefaultSocks5CommandResponse;
import io.netty.handler.codec.socksx.v5.Socks5AddressType;
import io.netty.handler.codec.socksx.v5.Socks5CommandRequestDecoder;
import io.netty.handler.codec.socksx.v5.Socks5CommandStatus;
import org.hk.proxy.common.constant.CommonKey;
import org.hk.proxy.common.constant.ResultCode;
import org.hk.proxy.protocol.ProxyMessage;
import org.hk.proxy.server.channel.ProxyChannelManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

public class ServerChannelHandler extends SimpleChannelInboundHandler<ProxyMessage> {
    private static final Logger log = LoggerFactory.getLogger(ServerChannelHandler.class);
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, ProxyMessage msg){
        switch (msg.getType()){
            case ProxyMessage.C_TYPE_AUTH:
                handlerAuthMsg(ctx,msg);
                break;
            case ProxyMessage.TYPE_CONNECT:
                handlerConnectMsg(ctx,msg);
                break;
            case ProxyMessage.TYPE_DISCONNECT:
                handleDisconnectMessage(ctx, msg);
                break;
            case ProxyMessage.P_TYPE_TRANSFER:
                handleTransferMessage(ctx, msg);
                break;
            case ProxyMessage.HEART_PACKAGE:
                handlerHeatPackage(ctx);
                break;
            case ProxyMessage.PROXY_SOCKET5_CONNECT:
                handlerSocket5Connect(ctx);
                break;
            default:
                break;
        }
    }

    private void handlerSocket5Connect(ChannelHandlerContext ctx){
        log.info("连接目标服务器成功,channel:{}", ctx.channel());
        Channel realChannel = ctx.channel().attr(CommonKey.NEXT_CHANNEL).get();
        Socks5AddressType socks5AddressType = realChannel.attr(CommonKey.Socks5AddressType).get();
        DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS,socks5AddressType);
        realChannel.pipeline().remove(Socks5CommandRequestDecoder.class);
        realChannel.pipeline().remove(Socks5CommandRequestInboundHandler.class);
        realChannel.pipeline().addLast(new NextChannelExchangeHandler());
        realChannel.config().setOption(ChannelOption.AUTO_READ,true);
        ctx.channel().config().setOption(ChannelOption.AUTO_READ,true);
        realChannel.writeAndFlush(commandResponse);
    }

    private void handlerHeatPackage(ChannelHandlerContext ctx){
        ProxyMessage proxyMessage = new ProxyMessage();
        proxyMessage.setType(ProxyMessage.HEART_PACKAGE);
        ctx.channel().writeAndFlush(proxyMessage);
    }
    /**
     * @Author LiuXu
     * @Description //客户端连接服务端校验逻辑
     * @Date 0:23 2022/11/21
     * @Param
     * @return
     **/
    private void handlerAuthMsg(ChannelHandlerContext ctx, ProxyMessage msg){
        String clientKey = msg.getUri();
        log.info("clientKey:{}",clientKey);
        if(StrUtil.isNotBlank(clientKey)){
            if(ProxyChannelManager.auth(clientKey)){
                ProxyChannelManager.addCmdProxy(clientKey,ctx.channel());
                ProxyMessage retMsg = new ProxyMessage();
                retMsg.setData(ByteUtil.intToBytes(ResultCode.SUCCESS.getCode()));
                retMsg.setType(ProxyMessage.C_TYPE_AUTH);
                ctx.channel().writeAndFlush(retMsg);
            }
        }else {
            log.error("clientKey is null.{}",ctx.channel());
            ctx.channel().close();
        }
    }
    /**
     * 数据转运
     */
    private void handleTransferMessage(ChannelHandlerContext ctx,ProxyMessage msg){
        Channel userChannel = ctx.channel().attr(CommonKey.PROXY_USER).get();
        Channel clientChannel = ctx.channel().attr(CommonKey.NEXT_CHANNEL).get();
        ByteBuf buf = ctx.alloc().buffer(msg.getData().length);
        buf.writeBytes(msg.getData());
        if(userChannel!=null){
            userChannel.writeAndFlush(buf);
        }else if(clientChannel!=null){
            clientChannel.writeAndFlush(buf);
        }
    }
    private void handlerConnectMsg(ChannelHandlerContext ctx,ProxyMessage msg){
        if(msg==null){
            ctx.channel().close();
        }
        Channel userChannel =  ProxyChannelManager.currentUserChannel.get(msg.getSerialNumber());
        if(userChannel!=null){
            ctx.channel().attr(CommonKey.PROXY_USER).set(userChannel);
            userChannel.attr(CommonKey.PROXY_USER).set(ctx.channel());
            userChannel.config().setOption(ChannelOption.AUTO_READ,true);
        }
    }
    private void handleDisconnectMessage(ChannelHandlerContext ctx, ProxyMessage proxyMessage) {
        log.info("real server inactive!");
        Channel userChannel = ctx.channel().attr(CommonKey.PROXY_USER).getAndSet(null);
        if(userChannel!=null){
            // 数据发送完成后再关闭连接，解决http1.0数据传输问题
            userChannel.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
        }
        Channel proxyChannel = ctx.channel().attr(CommonKey.NEXT_CHANNEL).getAndSet(null);
        if(proxyChannel!=null){
            log.error("连接目标服务器失败,channel:{}", proxyChannel);
            DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, proxyChannel.attr(CommonKey.Socks5AddressType).get());
            proxyChannel.writeAndFlush(commandResponse);
            proxyChannel.close();
        }
        return;
    }
    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        ProxyChannelManager.removeCmdProxy(ctx.channel());
        super.channelInactive(ctx);
    }
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.error("server channel error:{}",cause);
        if(cause instanceof IOException){
            channelInactive(ctx);
        }else {
            super.exceptionCaught(ctx,cause);
        }
    }
    @Override
    public void channelWritabilityChanged(ChannelHandlerContext ctx) throws Exception {
        Channel userChannel = ctx.channel().attr(CommonKey.PROXY_USER).get();
        if (userChannel != null) {
            userChannel.config().setOption(ChannelOption.AUTO_READ, ctx.channel().isWritable());
        }
        super.channelWritabilityChanged(ctx);
    }
}
